#include "Shader.h"
#include "Basic/Stream.h"
#include "Basic/FileSystem.h"
#include "Windows/GxDebug.h"

namespace FinalRendering
{
	namespace Engine
	{
		RefPtr<MetaLexer> Program::lexer;

		class ShaderException : public Exception
		{
		public:
			ShaderException(const String & txt)
				:Exception(txt)
			{
			}
		};

		class UniformNotFoundException : public Exception
		{
		public:
			UniformNotFoundException(const String & txt)
				:Exception(String(L"Uniform not found: ") + txt)
			{
			}
		};

		GPU_Shader::GPU_Shader(const String & fileName)
			: handle(0)
		{
			String ext = ExtractFileEXT(fileName);
			ext.Lowercase();
			if (ext == L"vs")
				handle = Shader::CreateShaderObjectARB(GL_VERTEX_SHADER);
			else
				handle = Shader::CreateShaderObjectARB(GL_FRAGMENT_SHADER);
			Shader::LoadShaderFromFile(handle, fileName);
			Shader::CompileShaderARB(handle);
			int succ;
			char errLog[4096];
			Shader::GetObjectParameterivARB(handle, GL_OBJECT_COMPILE_STATUS_ARB, &succ);
			if (succ == 0)
			{
				Shader::GetInfoLogARB(handle, sizeof(errLog), 0, errLog);
				throw ShaderException(String(L"Cannot compile shader \"") + fileName + L"\"\n-------------\n" +errLog);
			}
		}

		GPU_Shader::~GPU_Shader()
		{
			Shader::DeleteObjectARB(handle);
		}

		UInt GPU_Shader::GetHandle()
		{
			return handle;
		}

		void Program::Finalize()
		{
			lexer = 0;
		}

		class ProgramTokenTypeClass
		{
		public:
			enum _ProgramTokenType
			{
				Int, Float, String, Comma, LBrace, RBrace, LBracket, RBracket, Identifier, Equal, Colon
			};
			static GxLibBasic::String ToString(_ProgramTokenType type)
			{
				switch (type)
				{
				case Int:
					return L"int constant";
				case Float:
					return L"float constant";
				case String:
					return L"String";
				case Comma:
					return L"\',\'";
				case LBrace:
					return L"\'{\'";
				case RBrace:
					return L"\'}\'";
				case LBracket:
					return L"\'[\'";
				case RBracket:
					return L"\']\'";
				case Identifier:
					return L"identifier";
				case Equal:
					return L"\'=\'";
				case Colon:
					return L"\':\'";
				default:
					return L"null";
				}
			}
		};

		typedef ProgramTokenTypeClass::_ProgramTokenType ProgramTokenType;

		class ProgramParseException : public Exception
		{
		public:
			ProgramParseException(const String & syntaxErrorText, const String & fileName)
				: Exception(String(L"Program Parse Exception \"") + syntaxErrorText + L"\" in \""
				+ fileName + L"\"")
			{

			};
		};

		LexStream::LinkNode * ReadToken(LexStream::LinkNode *& node, ProgramTokenType type, const String & fileName)
		{
			if (!node || node->Data.TypeID != (int) type)
				throw ProgramParseException(ProgramTokenTypeClass::ToString(type) + L" expected at " + node->Data.Position,
				fileName);
			LexStream::LinkNode * current = node;
			node = node->GetNext();
			return current;
		}

		void Program::ParseShaderNames(RefPtr<Technique> & tech, LexStream::LinkNode *& node)
		{
			String path = ExtractFilePath(fileName);
			LexStream::LinkNode * token = ReadToken(node, ProgramTokenTypeClass::Identifier, fileName);
			token = ReadToken(node, ProgramTokenTypeClass::LBrace, fileName);
			while (node && node->Data.TypeID != (int)ProgramTokenTypeClass::RBrace)
			{
				token = ReadToken(node, ProgramTokenTypeClass::String, fileName);
				if (token->Data.Str.GetLength() > 2)
				{
					RefPtr<GPU_Shader> shader = new GPU_Shader(path + token->Data.Str.Mid(1, token->Data.Str.GetLength()-2));
					tech->shaders.Add(shader);
				}
			}
			ReadToken(node, ProgramTokenTypeClass::RBrace, fileName);
		}

		void Program::ParseAutoUniforms(RefPtr<Technique> & tech, LexStream::LinkNode *& node)
		{
			LexStream::LinkNode * token = ReadToken(node, ProgramTokenTypeClass::Identifier, fileName);
			token = ReadToken(node, ProgramTokenTypeClass::LBrace, fileName);
			while (node && node->Data.TypeID != (int)ProgramTokenTypeClass::RBrace)
			{
				token = ReadToken(node, ProgramTokenTypeClass::Identifier, fileName);
				tech->autoUniforms.Add(token->Data.Str);
			}
			ReadToken(node, ProgramTokenTypeClass::RBrace, fileName);
		}

		void Program::ParseState(RefPtr<Technique> & tech, LexStream::LinkNode *& node)
		{
			LexStream::LinkNode * token = ReadToken(node, ProgramTokenTypeClass::Identifier, fileName);
			token = ReadToken(node, ProgramTokenTypeClass::LBrace, fileName);
			while (node && node->Data.TypeID != (int)ProgramTokenTypeClass::RBrace)
			{
				token = ReadToken(node, ProgramTokenTypeClass::Identifier, fileName);
				token->Data.Str.Lowercase();
				if (token->Data.Str == L"blend")
				{
					token = ReadToken(node, ProgramTokenTypeClass::Colon, fileName);
					token = ReadToken(node, ProgramTokenTypeClass::Identifier, fileName);
					String str = token->Data.Str.Lowercase();
					tech->blendFunc = Technique::None;
					if (str == L"none")
						tech->blendFunc = Technique::None;
					else if (str == L"alphablend")
						tech->blendFunc = Technique::AlphaBlend;
					else if (str == L"add")
						tech->blendFunc = Technique::Add;
				}
				else if (token->Data.Str == L"process")
				{
					token = ReadToken(node, ProgramTokenTypeClass::Colon, fileName);
					token = ReadToken(node, ProgramTokenTypeClass::Identifier, fileName);
					String str = token->Data.Str.Lowercase();
					tech->process = Technique::Normal;
					if (str == L"normal")
						tech->process = Technique::Normal;
					else if (str == L"reflect")
						tech->process = Technique::Reflect;
					else if (str == L"reflectandrefract")
						tech->process = Technique::ReflectAndRefract;
				}
			}
			ReadToken(node, ProgramTokenTypeClass::RBrace, fileName);
		}

		void Program::ParseProgramConfig(const String & config)
		{
			bool uniform = false;
			if (!lexer)
			{
				String profile =
					L"Int = {-?\\d+}\n"\
					L"Float = {-?\\d+(.\\d*)?}\n"\
					L"String = {\"([^\\\\\"]|\\\\\\.)*\"}\n"\
					L"Comma = {,}\n"\
					L"LBrace = {{}\n"\
					L"RBrace = {}}\n"\
					L"LBracket = {\\[}\n"\
					L"RBracket = {\\]}\n"\
					L"Identifier = {[a-zA-Z_]\\w*}\n"\
					L"Equal = {=}\n"\
					L"Colon = {:}\n"\
					L"#WhiteSpace = {\\s+}\n"\
					L"#SingleLineComment = {//[^\\n]*\\n}\n"\
					L"#MultiLineComment = {/\\*([^*]|\\*[^/])*\\*/}\n";
				lexer = new MetaLexer(profile);
			}
			LexStream lexStream;
			lexer->Parse(config, lexStream);
			LexStream::LinkNode * node = lexStream.Begin();
			defaultTechnique = reflectTechnique = -1;
			while (node)
			{
				node->Data.Str.Lowercase();
				if (node->Data.Str == L"technique")
					ParseTechnique(node);
				else
					throw ProgramParseException(L"Illegal section name. Expected \"technique\".", fileName);
			}
			if (defaultTechnique != -1)
				curTech = techniques[defaultTechnique].operator->();
			else
				throw ProgramParseException(String(L"Default technique not defined."), fileName);
		}

		void Program::ParseTechnique(LexStream::LinkNode *& node)
		{
			LexStream::LinkNode * token = ReadToken(node, ProgramTokenTypeClass::Identifier, fileName);
			token = ReadToken(node, ProgramTokenTypeClass::Identifier, fileName);
			RefPtr<Technique> tech = new Technique();
			tech->name = token->Data.Str;
			token = ReadToken(node, ProgramTokenTypeClass::LBrace, fileName);
			while (node && node->Data.TypeID != (int)ProgramTokenTypeClass::RBrace)
			{
				node->Data.Str.Lowercase();
				if (node->Data.Str == L"shaders")
					ParseShaderNames(tech, node);
				else if (node->Data.Str == L"autouniforms")
					ParseAutoUniforms(tech, node);
				else if (node->Data.Str == L"state")
					ParseState(tech, node);
				else
					throw ProgramParseException(L"Illegal section name. Expected \"shaders\" \"autouniforms\" or \"state\".", fileName);
			}
			ReadToken(node, ProgramTokenTypeClass::RBrace, fileName);
			techniques.Add(tech);
			tech->name.Lowercase();
			if (tech->name == L"default")
				defaultTechnique = techniques.Count()-1;
			else if (tech->name == L"reflect")
				reflectTechnique = techniques.Count()-1;
			try
			{
				tech->Link();
			}
			catch (const ShaderException & ex)
			{
				ShaderException n(String(L"Cannot link program \"") + fileName + L"\", " + ex.Text);
				throw n;
			}
		}

		Program::Technique::Technique()
			: process(Normal), blendFunc(None), handle(-1)
		{
		}

		bool Program::IsReflective()
		{
			return (reflectTechnique != -1);
		}

		Program::Technique::~Technique()
		{
			if (handle != -1)
				Shader::DeleteObjectARB(handle);
		}

		int Program::GetAutoUniformCount()
		{
			return curTech->autoUniforms.Count();
		}

		String Program::GetAutoUniform(int i)
		{
			return curTech->autoUniforms[i];
		}

		void Program::Technique::Link()
		{
			handle = Shader::CreateProgramObjectARB();
			for (int i=0; i<shaders.Count(); i++)
				Shader::AttachObjectARB(handle, shaders[i]->GetHandle());
			int linked = 0;
			char errLog[4096];
			memset(errLog, 0, sizeof(errLog));

			Shader::LinkProgramARB(handle);
			Shader::GetObjectParameterivARB(handle, GL_OBJECT_LINK_STATUS_ARB, &linked);

			Shader::GetInfoLogARB(handle, sizeof(errLog), NULL, errLog);

			if (linked == 0)
			{
				Shader::GetInfoLogARB(handle, sizeof(errLog), NULL, errLog);
  				throw ShaderException("technique \"" + name +"\"\n"+errLog);
			}
		}

		Program::Program(const String & _fileName)
		{
			fileName = _fileName;
			String text = ReadTextFile(fileName);
			curTech = 0;
			ParseProgramConfig(text);
			
		}

		bool Program::Validate()
		{
			int validate = 0;
			Shader::ValidateProgramARB(curTech->handle);
			Shader::GetObjectParameterivARB(curTech->handle, GL_OBJECT_VALIDATE_STATUS_ARB, &validate);
			char errLog[4096];
			Shader::GetInfoLogARB(curTech->handle, sizeof(errLog), NULL, errLog);
			if (strlen(errLog))
				GxWinSystem::Debug::WriteLine(errLog);
			if (validate == 0)
			{
				GxWinSystem::Debug::WriteLine(String(L"Program validation failed. \"") + fileName + L"\"\n"+errLog);
				return false;
			}
			return true;
		}

		void Program::BindFragDataLocation(const String & outputVar, int drawBufferID)
		{
			Shader::BindFragDataLocation(curTech->handle, drawBufferID, outputVar.ToMBString());
		}

		Program::~Program()
		{
			techniques.Clear(true);
		}

		Program::Technique::_Process Program::GetProcess()
		{
			return curTech->process;
		}

		void Program::SetTechnique(TechniqueType tech)
		{
			switch (tech)
			{
			case NormalTechnique:
				if (defaultTechnique >=0 && defaultTechnique < techniques.Count())
					curTech = techniques[defaultTechnique].operator->();
				break;
			case ReflectTechnique:
				if (reflectTechnique >=0 && reflectTechnique < techniques.Count())
					curTech = techniques[reflectTechnique].operator->();
				else
					throw ShaderException(String(L"Reflect technique not defined in program \"") + fileName + L"\"");
				break;
			}
		}

		void Program::Bind()
		{
			Shader::UseProgramObjectARB(curTech->handle);
			if (curTech->blendFunc == Technique::None)
				glDisable(GL_BLEND);
			else
			{
				glEnable(GL_BLEND);
				if (curTech->blendFunc == Technique::AlphaBlend)
					glBlendFunc(GL_SRC_ALPHA, GL_ONE_MINUS_SRC_ALPHA);
				else if (curTech->blendFunc == Technique::Add)
					glBlendFunc(GL_ONE, GL_ONE);
			}
		}

		void Program::UnBind()
		{
			Shader::UseProgramObjectARB(0);
		}

		UInt Program::GetHandle()
		{
			return curTech->handle;
		}

		int Program::GetUniformLocation(const String & name)
		{
			return Shader::GetUniformLocationARB(curTech->handle, name.ToMBString());
		}

		void Program::SetUniform(const String &name, float val)
		{
			int loc = Shader::GetUniformLocationARB(curTech->handle, name.ToMBString());
	#ifdef _DEBUG
			if (loc == -1)
			{
				//throw UniformNotFoundException(name);
			}
	#endif
			Shader::Uniform1fARB(loc, val);

		}

		void Program::SetUniform(const String &name, int val)
		{
			int loc = Shader::GetUniformLocationARB(curTech->handle, name.ToMBString());
	#ifdef _DEBUG
			if (loc == -1)
			{
				//throw UniformNotFoundException(name);
			}
	#endif
			Shader::Uniform1iARB(loc, val);

		}

		void Program::SetUniform(const String &name, const vec3 & val)
		{
			int loc = Shader::GetUniformLocationARB(curTech->handle, name.ToMBString());
	#ifdef _DEBUG
			if (loc == -1)
			{
				//throw UniformNotFoundException(name);
			}
	#endif
			Shader::Uniform3fARB(loc, val.x, val.y, val.z);

		}

		void Program::SetUniform(const String &name, const vec4 & val)
		{
			int loc = Shader::GetUniformLocationARB(curTech->handle, name.ToMBString());
	#ifdef _DEBUG
			if (loc == -1)
			{
				//throw UniformNotFoundException(name);
			}
	#endif
			Shader::Uniform4fARB(loc, val.x,val.y,val.z,val.w);
		}

		void Program::SetUniform(int loc, int val)
		{
			Shader::Uniform1iARB(loc, val);
		}

		void Program::SetUniform(int loc, float val)
		{
			Shader::Uniform1fARB(loc, val);
		}

		void Program::SetUniform(int loc, const vec3 & val)
		{
			Shader::Uniform3fARB(loc, val.x,val.y,val.z);
		}

		void Program::SetUniform(int loc, const vec4 & val)
		{
			Shader::Uniform4fARB(loc, val.x,val.y,val.z,val.w);
		}
	}
}